"""
MyLangExtract 核心数据模型

定义了提取结果、示例数据等核心数据结构
"""

from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional
import json


@dataclass
class Extraction:
    """表示一个提取结果
    
    Attributes:
        extraction_class: 提取类别（如 "person", "location", "organization" 等）
        extraction_text: 从原文中提取的文本
        attributes: 附加属性字典，用于存储额外信息
        start_pos: 在原文中的起始位置（可选）
        end_pos: 在原文中的结束位置（可选）
    """
    extraction_class: str
    extraction_text: str
    attributes: Optional[Dict[str, Any]] = None
    start_pos: Optional[int] = None
    end_pos: Optional[int] = None
    
    def __post_init__(self):
        if self.attributes is None:
            self.attributes = {}
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        result = {
            "extraction_class": self.extraction_class,
            "extraction_text": self.extraction_text,
            "attributes": self.attributes or {}
        }
        if self.start_pos is not None:
            result["start_pos"] = self.start_pos
        if self.end_pos is not None:
            result["end_pos"] = self.end_pos
        return result
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Extraction":
        """从字典创建 Extraction 对象"""
        return cls(
            extraction_class=data["extraction_class"],
            extraction_text=data["extraction_text"],
            attributes=data.get("attributes"),
            start_pos=data.get("start_pos"),
            end_pos=data.get("end_pos")
        )


@dataclass
class ExampleData:
    """表示一个训练示例
    
    Attributes:
        text: 示例文本
        extractions: 从该文本中提取的结果列表
    """
    text: str
    extractions: List[Extraction] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            "text": self.text,
            "extractions": [ext.to_dict() for ext in self.extractions]
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ExampleData":
        """从字典创建 ExampleData 对象"""
        extractions = [
            Extraction.from_dict(ext_data) 
            for ext_data in data.get("extractions", [])
        ]
        return cls(
            text=data["text"],
            extractions=extractions
        )


@dataclass
class ExtractResult:
    """表示提取操作的完整结果
    
    Attributes:
        text: 原始输入文本
        extractions: 提取结果列表
        metadata: 元数据，包含模型信息、处理时间等
        success: 是否成功提取
        error_message: 错误信息（如果有）
    """
    text: str
    extractions: List[Extraction] = field(default_factory=list)
    metadata: Optional[Dict[str, Any]] = None
    success: bool = True
    error_message: Optional[str] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        result = {
            "text": self.text,
            "extractions": [ext.to_dict() for ext in self.extractions],
            "metadata": self.metadata or {},
            "success": self.success
        }
        if self.error_message:
            result["error_message"] = self.error_message
        return result
    
    def to_json(self, indent: int = 2) -> str:
        """转换为 JSON 字符串"""
        return json.dumps(self.to_dict(), ensure_ascii=False, indent=indent)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ExtractResult":
        """从字典创建 ExtractResult 对象"""
        extractions = [
            Extraction.from_dict(ext_data) 
            for ext_data in data.get("extractions", [])
        ]
        return cls(
            text=data["text"],
            extractions=extractions,
            metadata=data.get("metadata"),
            success=data.get("success", True),
            error_message=data.get("error_message")
        )
    
    @classmethod
    def from_json(cls, json_str: str) -> "ExtractResult":
        """从 JSON 字符串创建 ExtractResult 对象"""
        data = json.loads(json_str)
        return cls.from_dict(data)
    
    def add_extraction(self, extraction: Extraction):
        """添加一个提取结果"""
        self.extractions.append(extraction)
    
    def get_extractions_by_class(self, extraction_class: str) -> List[Extraction]:
        """根据类别获取提取结果"""
        return [ext for ext in self.extractions if ext.extraction_class == extraction_class]
    
    def get_all_classes(self) -> List[str]:
        """获取所有提取类别"""
        return list(set(ext.extraction_class for ext in self.extractions))


# 用于 Function Calling 的 JSON Schema 生成
def generate_extraction_schema(examples: List[ExampleData]) -> Dict[str, Any]:
    """根据示例生成 JSON Schema，用于 Function Calling
    
    Args:
        examples: 示例数据列表
        
    Returns:
        符合 OpenAI Function Calling 格式的 JSON Schema
    """
    # 收集所有的提取类别和属性
    classes = set()
    attributes_by_class = {}
    
    for example in examples:
        for extraction in example.extractions:
            classes.add(extraction.extraction_class)
            
            if extraction.extraction_class not in attributes_by_class:
                attributes_by_class[extraction.extraction_class] = set()
            
            if extraction.attributes:
                for attr_key in extraction.attributes.keys():
                    attributes_by_class[extraction.extraction_class].add(attr_key)
    
    # 生成 schema
    schema = {
        "type": "object",
        "properties": {
            "extractions": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "extraction_class": {
                            "type": "string",
                            "enum": list(classes),
                            "description": "提取结果的类别"
                        },
                        "extraction_text": {
                            "type": "string",
                            "description": "从原文中提取的文本"
                        },
                        "attributes": {
                            "type": "object",
                            "description": "附加属性",
                            "additionalProperties": {
                                "type": "string"
                            }
                        }
                    },
                    "required": ["extraction_class", "extraction_text"],
                    "additionalProperties": False
                }
            }
        },
        "required": ["extractions"],
        "additionalProperties": False
    }
    
    return schema
